import torchvision
from torchvision.models.detection import fasterrcnn_resnet50_fpn,FasterRCNN
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models import densenet121
from torchvision.models.detection.rpn import AnchorGenerator
import torch
# 必须使用该方法下载模型，然后加载
from flyai.utils import remote_helper
from efficientnet_pytorch import EfficientNet

def get_model(num_classes = 2):
    model = fasterrcnn_resnet50_fpn(pretrained=True)
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
    return model

def get_model_denmodel(num_classes = 2):
    #backbone = densenet121(pretrained=True).features
    denmodel = densenet121(pretrained=False)
    path = remote_helper.get_remote_data('https://www.flyai.com/m/densenet121-a639ec97.pth')
    from torchvision.models.densenet import _load_state_dict 
    _load_state_dict(denmodel,path,progress=True)
    backbone = denmodel.features
    backbone.out_channels = 1024
    anchor_generator = AnchorGenerator(sizes=(16,32,64,128,256),aspect_ratios=(0.5,1,1.5,2))
    roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],output_size=7,sampling_ratio=2)
    model = FasterRCNN(backbone,num_classes=num_classes,rpn_anchor_generator=anchor_generator,box_roi_pool=roi_pooler)
    return model

def get_model_2(num_classes=2):
    # Some basic setup:
    # Setup detectron2 logger
    import detectron2
    from detectron2.utils.logger import setup_logger
    setup_logger()

    # import some common libraries
    import numpy as np
    import os, json, cv2, random
    from google.colab.patches import cv2_imshow

    # import some common detectron2 utilities
    from detectron2 import model_zoo
    from detectron2.engine import DefaultPredictor
    from detectron2.config import get_cfg
    from detectron2.utils.visualizer import Visualizer
    from detectron2.data import MetadataCatalog, DatasetCatalog

    
    #visualize training data
    my_dataset_train_metadata = MetadataCatalog.get("my_dataset_train")
    dataset_dicts = DatasetCatalog.get("my_dataset_train")
    import random
    from detectron2.utils.visualizer import Visualizer
    for d in random.sample(dataset_dicts, 3):
        img = cv2.imread(d["file_name"])
        visualizer = Visualizer(img[:, :, ::-1], metadata=my_dataset_train_metadata, scale=0.5)
        vis = visualizer.draw_dataset_dict(d)
        cv2_imshow(vis.get_image()[:, :, ::-1])


    cfg = get_cfg()
    cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml"))
    cfg.DATASETS.TRAIN = ("my_dataset_train",)
    cfg.DATASETS.TEST = ("my_dataset_val",)
    cfg.DATALOADER.NUM_WORKERS = 4
    cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml")  # Let training initialize from model zoo
    cfg.SOLVER.IMS_PER_BATCH = 4
    cfg.SOLVER.BASE_LR = 0.001
    cfg.SOLVER.WARMUP_ITERS = 1000
    cfg.SOLVER.MAX_ITER = 1500 #adjust up if val mAP is still rising, adjust down if overfit
    cfg.SOLVER.STEPS = (1000, 1500)
    cfg.SOLVER.GAMMA = 0.05
    cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 64
    cfg.MODEL.ROI_HEADS.NUM_CLASSES = 4
    cfg.TEST.EVAL_PERIOD = 500



if __name__ == "__main__":
    get_model()